import os

from io import StringIO
import pandas

import glob

df1s = []
df2s = []

for job_type in ["aliyun", "aws"]:
  for filename in glob.glob(f"download/{job_type}*.txt"):
    with open(filename, 'r') as file:
      readbulk = file.read()
      if readbulk.find("Exception") >= 0:
        continue
      if readbulk.find("Oops") >= 0:
        continue
      tbls = readbulk.split("Checking")
      tbls = [tbl.split("Aborting.")[-1] for tbl in tbls]

      try:
        df1 = pandas.read_fwf(StringIO(tbls[1][tbls[1].find("\n"):]))
        df2 = pandas.read_fwf(StringIO(tbls[2][tbls[2].find("\n"):]))
        df1.iloc[0, 0] = "timestamp"
        df2.iloc[0, 0] = "timestamp"
        df1 = df1.rename(columns=df1.iloc[0]).iloc[1:]
        df2 = df2.rename(columns=df2.iloc[0]).iloc[1:]
        df1s.append(df1)
        df2s.append(df2)
      except Exception:
        pass

  try:
    dfa = pandas.concat(df1s, axis=0, sort=True).reset_index(drop=True)
  except Exception:
    import pdb
    pdb.set_trace()

  try:
    dfb = pandas.concat(df2s, axis=0, sort=True).reset_index(drop=True)
  except Exception:
    import pdb
    pdb.set_trace()

  os.makedirs('cpu_usage', exist_ok=True)

  import matplotlib
  matplotlib.use("Agg")
  import matplotlib.pyplot as plt

  plt.rcParams['figure.figsize'] = 8, 5
  plt.rcParams['font.size'] = 10
  plt.rcParams['legend.fontsize'] = 10
  plt.rcParams['xtick.labelsize'] = 10
  plt.rcParams['ytick.labelsize'] = 10

  for seldf in [dfa, dfb]:
    for tag, subdf in (seldf.groupby('tag') if 'tag' in seldf else seldf.groupby('instance_name')):
      dti = pandas.to_datetime(subdf['timestamp'], format='%Y%m%d:%H:%M:%S')
      try:
        plt.subplot(211)
        plt.plot(dti, subdf['cpu_usage'].astype(float), 'r.-', lw=0.2, markersize=1)
        plt.xticks(rotation=20)
        plt.xlabel('timestamp')
        plt.ylabel('cpu_usage')
        plt.title(tag)
        plt.subplot(212)
        plt.plot(dti,
                 subdf['cpu_usage'].astype(float).rolling(10).mean(),
                 'r.-',
                 lw=0.2,
                 markersize=1)
        plt.xticks(rotation=20)
        plt.xlabel('timestamp')
        plt.ylabel('cpu_usage')
        plt.title(tag)
        plt.title(f'{tag}_movingwindow')
        plt.tight_layout()
        plt.savefig(f"cpu_usage/{tag}.png")
        plt.close()
      except Exception:
        pass
